// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/file/file_systems.h>
#include <pollux/connectors/hive/table_handle.h>
#include <pollux/connectors/hive/iceberg/iceberg_delete_file.h>
#include <pollux/connectors/hive/iceberg/iceberg_metadata_columns.h>
#include <pollux/connectors/hive/iceberg/iceberg_split.h>
#include <pollux/connectors/hive/iceberg/iceberg_split_reader.h>
#include <pollux/testing/dwio/data_set_builder.h>
#include <pollux/dwio/dwrf/register_dwrf_reader.h>
#include <pollux/dwio/dwrf/writer/writer.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/testing/vector/vector_test_base.h>

#include <melon/benchmark.h>
#include <melon/init/init.h>

namespace kumo::pollux::iceberg::reader::test {

constexpr uint32_t kNumRowsPerBatch = 20000;
constexpr uint32_t kNumBatches = 50;
constexpr uint32_t kNumRowsPerRowGroup = 10000;
constexpr double kFilterErrorMargin = 0.2;

class IcebergSplitReaderBenchmark {
 public:
  explicit IcebergSplitReaderBenchmark() {
    rootPool_ =
        memory::memoryManager()->addRootPool("IcebergSplitReaderBenchmark");
    leafPool_ = rootPool_->addLeafChild("IcebergSplitReaderBenchmark");
    dataSetBuilder_ =
        std::make_unique<kumo::pollux::test::DataSetBuilder>(*leafPool_, 0);
    filesystems::registerLocalFileSystem();
    dwrf::registerDwrfReaderFactory();
  }

  ~IcebergSplitReaderBenchmark() {}

  void writeToFile(const std::vector<RowVectorPtr>& batches);

  void writeToPositionDeleteFile(
      const std::string& filePath,
      const std::vector<RowVectorPtr>& vectors);

  dwio::common::FilterSpec createFilterSpec(
      const std::string& columnName,
      float startPct,
      float selectPct,
      const TypePtr& type,
      bool isForRowGroupSkip,
      bool allowNulls);

  std::shared_ptr<dwio::common::ScanSpec> createScanSpec(
      const std::vector<RowVectorPtr>& batches,
      RowTypePtr& rowType,
      const std::vector<dwio::common::FilterSpec>& filterSpecs,
      std::vector<uint64_t>& hitRows,
      std::unordered_map<
          kumo::pollux::common::Subfield,
          std::unique_ptr<kumo::pollux::common::Filter>>& filters);

  int read(
      const RowTypePtr& rowType,
      uint32_t nextSize,
      std::unique_ptr<connector::hive::iceberg::IcebergSplitReader>
          icebergSplitReader);

  void readSingleColumn(
      const std::string& columnName,
      const TypePtr& type,
      float startPct,
      float selectPct,
      float deleteRate,
      uint32_t nextSize);

  std::vector<std::shared_ptr<connector::hive::HiveConnectorSplit>>
  createIcebergSplitsWithPositionalDelete(
      int32_t deleteRowsPercentage,
      int32_t deleteFilesCount);

  std::vector<std::string> listFiles(const std::string& dirPath);

  std::shared_ptr<connector::hive::HiveConnectorSplit> makeIcebergSplit(
      const std::string& dataFilePath,
      const std::vector<connector::hive::iceberg::IcebergDeleteFile>&
          deleteFiles = {});

  std::vector<int64_t> makeRandomDeleteRows(int32_t deleteRowsCount);

  std::vector<int64_t> makeSequenceRows(int32_t maxRowNumber);

  std::string writePositionDeleteFile(
      const std::string& dataFilePath,
      int64_t numDeleteRows);

 private:
  const std::string fileName_ = "test.data";
  const std::shared_ptr<exec::test::TempDirectoryPath> fileFolder_ =
      exec::test::TempDirectoryPath::create();
  const std::shared_ptr<exec::test::TempDirectoryPath> deleteFileFolder_ =
      exec::test::TempDirectoryPath::create();

  std::unique_ptr<kumo::pollux::test::DataSetBuilder> dataSetBuilder_;
  std::shared_ptr<memory::MemoryPool> rootPool_;
  std::shared_ptr<memory::MemoryPool> leafPool_;
  std::unique_ptr<dwrf::Writer> writer_;
  dwio::common::RuntimeStatistics runtimeStats_;

  dwio::common::FileFormat fileFomat_{dwio::common::FileFormat::DWRF};
  const std::string kHiveConnectorId = "hive-iceberg";
};

void run(
    uint32_t,
    const std::string& columnName,
    const TypePtr& type,
    float filterRateX100,
    float deleteRateX100,
    uint32_t nextSize);

} // namespace kumo::pollux::iceberg::reader::test
